from elftools.elf.elffile import ELFFile
from elftools.dwarf.descriptions import describe_form_class
import subprocess
import argparse
import pylink
import json
import sys
import os
import re

VERSION = '1.0.0'

def decode_funcname(dwarfinfo, address):
    # Go over all DIEs in the DWARF information, looking for a subprogram
    # entry with an address range that includes the given address. Note that
    # this simplifies things by disregarding subprograms that may have
    # split address ranges.
    for CU in dwarfinfo.iter_CUs():
        for DIE in CU.iter_DIEs():
            try:
                if DIE.tag == 'DW_TAG_subprogram':
                    lowpc = DIE.attributes['DW_AT_low_pc'].value

                    # DWARF v4 in section 2.17 describes how to interpret the
                    # DW_AT_high_pc attribute based on the class of its form.
                    # For class 'address' it's taken as an absolute address
                    # (similarly to DW_AT_low_pc); for class 'constant', it's
                    # an offset from DW_AT_low_pc.
                    highpc_attr = DIE.attributes['DW_AT_high_pc']
                    highpc_attr_class = describe_form_class(highpc_attr.form)
                    if highpc_attr_class == 'address':
                        highpc = highpc_attr.value
                    elif highpc_attr_class == 'constant':
                        highpc = lowpc + highpc_attr.value
                    else:
                        print('Error: invalid DW_AT_high_pc class:',
                              highpc_attr_class)
                        continue

                    if lowpc <= address < highpc:
                        return DIE.attributes['DW_AT_name'].value
            except KeyError:
                continue
    return None

def decode_file_line(dwarfinfo, address):
    # Go over all the line programs in the DWARF information, looking for
    # one that describes the given address.
    for CU in dwarfinfo.iter_CUs():
        # First, look at line programs to find the file/line for the address
        lineprog = dwarfinfo.line_program_for_CU(CU)
        if lineprog is None:
            continue

        delta = 1 if lineprog.header.version < 5 else 0
        prevstate = None
        for entry in lineprog.get_entries():
            # We're interested in those entries where a new state is assigned
            if entry.state is None:
                continue
            # Looking for a range of addresses in two consecutive states that
            # contain the required address.
            if prevstate and prevstate.address <= address < entry.state.address:
                filename = lineprog['file_entry'][prevstate.file - delta].name
                line = prevstate.line
                return filename, line
            if entry.state.end_sequence:
                # For the state with `end_sequence`, `address` means the address
                # of the first byte after the target machine instruction
                # sequence and other information is meaningless. We clear
                # prevstate so that it's not used in the next iteration. Address
                # info is used in the above comparison to see if we need to use
                # the line information for the prevstate.
                prevstate = None
            else:
                prevstate = entry.state
    return None, None

def get_function_info_by_addr(dwarfinfo, addr, path):
    '''
    通过地址获取函数名、源码文件路径和函数名。并打印keil可跳转
    的字符串到Build Output窗口中
    '''
    patten = re.compile('^[A-Fa-f0-9]+$')
    patten0x = re.compile('0x[0-9a-fA-F]+')
        
    if patten.match(addr) or patten0x.match(addr):
        funcname = decode_funcname(dwarfinfo, int(addr, base=16))
        file, line = decode_file_line(dwarfinfo, int(addr, base=16))
        filename = file.decode()
    
        # print(file.decode(), '(', line, ')', ':', funcname.decode())
        if path == None or path == '':
            print(filename, '(', line, ')', ':', funcname.decode())
        else:
            # 再输入的路径中查找文件并输出结果
            result = []
            for root, lists, files in os.walk(path):
                for file in files:
                    if filename in file:
                        write = os.path.join(root, file)
                        # 对文件后缀进行筛选，只保留C语言和c++相关的文件
                        if write.endswith(tuple(['c', 'cc', 'cpp', 'h'])):
                            result.append(write)

            for subpath in result:
                print(subpath, '({:d})'.format(line), ':', funcname.decode())

def funtion_push_info(asmpath, funcname):
    '''
    通过函数名再asm文件中获取该函数的压栈信息

    asmpath: 通过elf文件生成的asm文件
    funcname: 函数名

    ret: 返回压栈时的地址偏移和压入栈的寄存器列表
    '''
    push_str = []

    # 打开asm文件
    with open(asmpath, 'r') as asmf:
        # 将文件数据按行读入并解析
        lines = asmf.readlines()
        match_index = False
        for index in range(len(lines)):
            # 查找函数
            if lines[index].strip() == funcname:
                cunter = 0
                while True:
                    if lines[index].find('PUSH') > 0:
                        push_str.append(lines[index])
                        match_index = True
                    # 包含局部变量的函数再压栈后会在栈里给局部变量
                    # 分配空间，具体的做法就是用SUB指令将SP指针的
                    # 减去某个值。所以，我们解析栈数据时先看到的是
                    # 局部变量的值。
                    elif lines[index].find('SUB') > 0:
                        push_str.append(lines[index])
                    else:
                        if match_index:
                            break
                    index = index + 1
                    cunter = cunter + 1
                    if (not match_index) and cunter >= 2:
                        break
                break
        asmf.close()

    if len(push_str) <= 0:
        return 0, 0
    
    # 获取函数地址
    funcaddr = push_str[0].strip().split(':')[0]
    funcaddr = int(funcaddr, base=16)
    # 对压栈信息进行解析
    reglist = []
    subw = 0
    for line in push_str:
        # 获取到压栈指令行，并解析出寄存器列表
        if line.find('PUSH') > 0:
            start = line.find('{')
            end = line.find('}')
            line = line[start+1:end]
            reglist.extend(line.split(','))
            continue
        sub_start = line.find('SUB')
        if sub_start > 0:
            # SUB指令行不一定就是对sp指针的操作
            line = line[sub_start:]
            splist = line.split(',')
            # 确保是对sp的操作，并解出偏移
            if splist[1] == 'sp':
                offset = line.split('#')[1].rstrip()
                if offset[0:2] == '0x':
                    subw = subw + int(offset, base=16)
                else:
                    subw = subw + int(offset, base=10)

    reg_num = 0
    for reg in reglist:
        if reg.find('-') > 0:
            tmp_list = reg.split('-')
            reg_num = reg_num + (int(tmp_list[1][1:]) - int(tmp_list[0][1:]) + 1)
        else:
            reg_num = reg_num + 1

    return (reg_num*4 + subw), funcaddr

def vector_from_asm(asmpath):
    '''
    从asm文件中获取向量表
    '''
    vector = []
    with open(asmpath, 'r') as asmf:
        lines = asmf.readlines()
        flag = False
        for line in lines:
            if flag:
                if line.strip() == '$t':
                    break
                if line.find('DCD') > 0:
                    tmp = line.split('DCD')[1].strip()
                    vector.append(int(tmp, base=10))
            else:
                if line.strip() == 'RESET':
                    flag = True
                    continue
        asmf.close()

    return vector

def stack_backtrace(jlink, asmpath, dwarfinfo):
    '''
    对栈进行回溯

    regs: [sp, lr, pc]
    '''
    # 获取向量表
    vector = vector_from_asm(asmpath)
    # sp, lr, pc, IPSR
    regs = link.register_read_multiple([13, 14, 15, 16])
    result = []
    # 通过pc获取当前函数名
    funcname = decode_funcname(dwarfinfo, regs[2])
    file, line = decode_file_line(dwarfinfo, regs[2])
    result.append([regs[2], file.decode(), line, funcname.decode()])

    sp = regs[0]
    ipsr = regs[3] & 0x1F
    # 对栈进行回溯
    while True:
        offset, funcaddr = funtion_push_info(asmpath, funcname.decode())
        if offset == 0:
            break
        # 如果ipsr大于0，说明当前pc在中断处理函数当中
        # ARM Cortex-M处理器在进入中断处理函数前会自动
        # 将R0 <- R1 <- R2 <- R3 <- R12 <- LR <- PC <- xPSR
        # 压入栈中，针对这一特点需要对这种情况特殊处理
        if ipsr > 0 and funcaddr == vector[ipsr]-1:
            sp = sp + offset + 32
            lr = jlink.memory_read32(sp-8, 1)[0]
            funcname = decode_funcname(dwarfinfo, lr)
            file, line = decode_file_line(dwarfinfo, lr)
        else:
            sp = sp + offset
            lr = jlink.memory_read32(sp-4, 1)[0]
            funcname = decode_funcname(dwarfinfo, lr - 1)
            file, line = decode_file_line(dwarfinfo, lr - 1)
        if not funcname:
            break
        result.append([lr-1, file.decode(), line, funcname.decode()])

    return result

if __name__ == '__main__':
    parse = argparse.ArgumentParser(description='根据地址定位函数名、文件名和行号')
    parse.add_argument('-v', '--version', action='version', version=VERSION, help='Show version and exit.')
    parse.add_argument('-addr', help='函数地址')
    parse.add_argument('-json', help='Json配置文件的路径')

    args = parse.parse_args()

    # 参数检查
    if args.json is None:
        parse.print_help()
        sys.exit(0)

    # 读取json文件获取配置
    config = None
    try:
        with open(args.json, 'r') as jf:
            config = json.load(jf)
            jf.close()
    except:
        parse.print_help()
        sys.exit(0)

    elfpath = config.get('Elfpath')
    projectpath = config.get('Projectpath')
    if not (elfpath and projectpath):
        parse.print_help()
        sys.exit(0)

    elfpath = projectpath + elfpath

    with open(elfpath, 'rb') as f:
        elffile = ELFFile(f)

        if not elffile.has_dwarf_info():
            print('file has no SWARF info')
        else:
            dwarfinfo = elffile.get_dwarf_info()
            mode = config.get('Mode')
            if not mode:
                sys.exit(0)
            if mode == 'funcname':
                if args.addr != None:
                    get_function_info_by_addr(dwarfinfo, args.addr, projectpath)
            
            if mode == 'stack':
                keilpath = config.get('Keilpath')
                device = config.get('Device')
                if not (keilpath and device):
                    sys.exit(0)
                cmd = [keilpath + 'fromelf.exe', '--text', '-c', '-o', projectpath + '\\frimwar.asm', elfpath]
                subprocess.Popen(cmd)
                
                # 获取jlin的sn号
                obj = subprocess.Popen(['JLink'], shell=True, stdout=subprocess.PIPE, stdin=subprocess.PIPE)
                # time.sleep(10)
                obj.stdin.write('q\n'.encode())
                out, err = obj.communicate()
            
                sn = ''
                for line in out.splitlines():
                    str = line.decode()
                    if str.find('S/N') >= 0:
                        sn = str.split(':')[1].rstrip()
                        break
                
                # 连接jlink
                link = pylink.JLink()
                link.open(sn)
                link.set_tif(pylink.enums.JLinkInterfaces.SWD)
                link.connect(device)

                ok = link.halt()

                asmpath = projectpath + '\\frimwar.asm'
                result = stack_backtrace(link, asmpath, dwarfinfo)
                out_result = []
                for item in result:
                    filename = item[1].split('\\')[-1].rstrip()
                    for root, lists, files in os.walk(projectpath):
                        for file in files:
                            if filename in file:
                                write = os.path.join(root, file)
                                # 对文件后缀进行筛选，只保留C语言和c++相关的文件
                                if write.endswith(tuple(['c', 'cc', 'cpp', 'h'])):
                                    out_result.append([write, item])

                # 对栈帧进行打印
                print('backtrace:')
                for subpath in out_result:
                    print(subpath[0], '({:d})'.format(subpath[1][2]), ':', subpath[1][3])

                # 程序继续运行
                link.restart()

                os.remove(asmpath)
